package com.xja.utils;

import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.xja.model.entity.Model;
import com.xja.service.IModelService;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Component;

import static com.xja.constant.ModelType.TONGUE_MODEL_TYPE;

/**
 * @author 蔡建国
 * Date  2025/5/10 21:48
 */
@Component
@RequiredArgsConstructor
public class ChatClientHolder {

    private final IModelService modelService;
    private final VectorStore vectorStore;
    private final ChatClient.Builder builder;

    // 普通对话模型客户端
    @Getter
    private ChatClient simpleChatClient;
    // 健康评测模型客户端
    @Getter
    private ChatClient evaluateChatClient;
    // 舌苔诊断模型客户端
    @Getter
    private ChatClient tongueChatClient;

    public ChatClient getTestChatClient() {
        return builder.defaultSystem("你是一个智能助手").build();
    }

    // 刷新所有的模型Client
    public void refreshAllChatClient() {
      this.refreshSimpleChatClient();
      this.refreshEvaluateChatClient();
      this.refreshTongueChatClient();
    }

    // 刷新健康评测模型客户端
    private void refreshEvaluateChatClient() {
        // 获取当前选择的对话模型的信息
        Model model = modelService.getCurrentEvaluateModel();
        this.evaluateChatClient = getChatClient(model);
    }

    // 刷新普通对话模型客户端
    private void refreshSimpleChatClient() {
        // 获取当前选择的对话模型的信息
        Model model = modelService.getCurrentSimpleChatModel();
        this.simpleChatClient = getChatClient(model);
    }

    // 刷新普通对话模型客户端
    private void refreshTongueChatClient() {
        // 获取当前选择的对话模型的信息
        Model model = modelService.getcurrentTongueModel();
        this.tongueChatClient = getChatClient(model);
    }

    private ChatClient getChatClient(Model model) {
        // 根据apikey创建DashScopeApi实例
        DashScopeApi dashScopeApi = DashScopeApi.builder().apiKey(model.getApiKey()).build();

        // RAG顾问
        Advisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                // 配置查询增强器
                .queryAugmenter(ContextualQueryAugmenter.builder()
                        // 允许空上下文（没检索到也能查询）
                        .allowEmptyContext(true)
                        .build())
                // 配置文档检索器
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .vectorStore(vectorStore)
                        // 相似度阈值
                        .similarityThreshold(model.getThreshold())
                        // 最大返回文档数量
                        .topK(model.getTopK())
                        .build())
                .build();

        // 返回Client
        return ChatClient.builder(DashScopeChatModel.builder()
                        .dashScopeApi(dashScopeApi)
                        .build())
                // 设置系统提示词
                .defaultSystem(model.getDefaultSystem())
                // 添加RAG顾问
                .defaultAdvisors(retrievalAugmentationAdvisor)
                // 模型参数
                .defaultOptions(DashScopeChatOptions.builder()
                        // 模型名称
                        .withModel(model.getModel())
                        // 是否开启多模态
                        .withMultiModel(TONGUE_MODEL_TYPE.equals(model.getType()))
                        // 模型温度
                        .withTemperature(model.getTemperature())
                        .build())
                .build();
    }
}